#pragma once

namespace MockItNow
{
#pragma warning(disable:4700)

	template<typename T>
	int FixVirtualFunction(int address)
	{
		// A bit dodgy, we need to look at the first couple of instructions
		// in the function to see if it is a virtual function.
		unsigned char* instructions = (unsigned char*)address;

		// These are the byte codes for
		// mov ecx, dword ptr [eax]			// 8b 01 ff
		// jmp [ecx + ???] or jmp [ecx]		// 60 or 20
		// The offset in the vtable will be the next byte after these instructions if the
		// jump instruction was 60
		if (*(instructions + 0) == 0x8b && *(instructions + 1) == 0x01 && *(instructions + 2) == 0xff)
		{
			instructions += 3;

			int offset = 0;
			if (*instructions == 0x60)
				offset = *(instructions + 1);

			T object;
			const int* vtable = (int*)(*((int*)&object));
			return *(vtable + (offset / 4));
		}

		return address;
	}

#pragma warning(default:4700)
}
